#include <cmath>
#include <vector>
#include <Eigen/Geometry>
#include <pcl/point_types.h>
#include <pcl/point_cloud.h>
#include <pcl/kdtree/impl/kdtree_flann.hpp>

// 自定义点类型，包含四元数
struct QuaternionPoint {
    float x, y, z, w;
    // 默认构造函数
    QuaternionPoint() {
        x = y = z = w = 0.0f;
    }

    // 带参数的构造函数
    QuaternionPoint(float x, float y, float z, float w) : x(x), y(y), z(z), w(w) {}
    EIGEN_MAKE_ALIGNED_OPERATOR_NEW     // 确保 Eigen 对齐
} EIGEN_ALIGN16;                    // 确保 16 字节对齐

// 注册自定义点类型
POINT_CLOUD_REGISTER_POINT_STRUCT(QuaternionPoint,
                                  (float, x, x)
                                  (float, y, y)
                                  (float, z, z)
                                  (float, w, w))

template<typename T>
void GenerateRandomPointCloudQuat(pcl::PointCloud<QuaternionPoint>& pointcloud, const size_t N)
{
    // Generating Random Quaternions
    pointcloud.points.resize(N);
    T theta, X, Y, Z, sinAng, cosAng, mag;
    for (size_t i = 0; i < N; i++)
    {
        theta = static_cast<T>(M_PI * (((double)rand()) / RAND_MAX));
        // Generate random value in [-1, 1]
        X   = static_cast<T>(2 * (((double)rand()) / RAND_MAX) - 1);
        Y   = static_cast<T>(2 * (((double)rand()) / RAND_MAX) - 1);
        Z   = static_cast<T>(2 * (((double)rand()) / RAND_MAX) - 1);
        mag = sqrt(X * X + Y * Y + Z * Z);
        X /= mag;
        Y /= mag;
        Z /= mag;
        cosAng = cos(theta / 2);
        sinAng = sin(theta / 2);
        pointcloud.points[i].w = cosAng;
        pointcloud.points[i].x = X * sinAng;
        pointcloud.points[i].y = Y * sinAng;
        pointcloud.points[i].z = Z * sinAng;
    }
}

// 自定义距离函数，计算四元数之间的角度距离
float QuaternionDistance(const QuaternionPoint& p1, const QuaternionPoint& p2) {
    Eigen::Quaternionf q1(p1.w, p1.x, p1.y, p1.z);
    Eigen::Quaternionf q2(p2.w, p2.x, p2.y, p2.z);
    float angle = q1.angularDistance(q2);
    return angle;
}

int main(int argc, const char** argv) {
    // ... 加载点云数据 ...
    pcl::PointCloud<QuaternionPoint>::Ptr cloud(new pcl::PointCloud<QuaternionPoint>);
    GenerateRandomPointCloudQuat<float>(*cloud, 100);
    // 创建KdTree，使用自定义距离函数
    pcl::KdTreeFLANN<QuaternionPoint> kdtree;
    kdtree.setInputCloud(cloud);
    QuaternionPoint query_point(0,0,0,1);
    // 进行最近邻搜索
    double radius = 0.2;
    std::vector<int> k_indices;
    std::vector<float> k_sqr_distances;
    kdtree.radiusSearch(query_point, radius, k_indices, k_sqr_distances);
    for (size_t i = 0; i < k_indices.size(); i++)
    {
        if (QuaternionDistance(query_point, cloud->points[k_indices[i]]) < 2 * radius) 
        {
            std::cout << "Index: " << k_indices[i] << " Point: (" << cloud->points[k_indices[i]].x << " " << cloud->points[k_indices[i]].y
            << " " << cloud->points[k_indices[i]].z << " "<< cloud->points[k_indices[i]].w << " " << " Distance: " << k_sqr_distances[i]
            << " AngleDistance: " << QuaternionDistance(query_point, cloud->points[k_indices[i]]) << std::endl;
        }

    }

    return 0;
}
